# -*- coding: UTF-8 -*-


#how to use
#python doMD.py <pdb name> <times(numer: 1 ns)>
#

import os
import sys
name = sys.argv[1]
time = float(sys.argv[2])
#创建 mdp
ion_list = ['; ions.mdp - used as input into grompp to generate ions.tpr\n', '; Parameters describing what to do, when to stop and what to save\n', 'integrator  = steep         ; Algorithm (steep = steepest descent minimization)\n', 'emtol       = 1000.0        ; Stop minimization when the maximum force < 1000.0 kJ/mol/nm\n', 'emstep      = 0.01          ; Minimization step size\n', 'nsteps      = 50000         ; Maximum number of (minimization) steps to perform\n', '\n', '; Parameters describing how to find the neighbors of each atom and how to calculate the interactions\n', 'nstlist         = 1         ; Frequency to update the neighbor list and long range forces\n', 'cutoff-scheme\t= Verlet    ; Buffered neighbor searching \n', 'ns_type         = grid      ; Method to determine neighbor list (simple, grid)\n', 'coulombtype     = cutoff    ; Treatment of long range electrostatic interactions\n', 'rcoulomb        = 1.0       ; Short-range electrostatic cut-off\n', 'rvdw            = 1.0       ; Short-range Van der Waals cut-off\n', 'pbc             = xyz       ; Periodic Boundary Conditions in all 3 dimensions\n']
ion = open("ions.mdp","w+")
ion.writelines(ion_list)
ion.close()
mini_list = ['; minim.mdp - used as input into grompp to generate em.tpr\n', '; Parameters describing what to do, when to stop and what to save\n', 'integrator  = steep         ; Algorithm (steep = steepest descent minimization)\n', 'emtol       = 1000.0        ; Stop minimization when the maximum force < 1000.0 kJ/mol/nm\n', 'emstep      = 0.01          ; Minimization step size\n', 'nsteps      = 50000         ; Maximum number of (minimization) steps to perform\n', '\n', '; Parameters describing how to find the neighbors of each atom and how to calculate the interactions\n', 'nstlist         = 1         ; Frequency to update the neighbor list and long range forces\n', 'cutoff-scheme   = Verlet    ; Buffered neighbor searching\n', 'ns_type         = grid      ; Method to determine neighbor list (simple, grid)\n', 'coulombtype     = PME       ; Treatment of long range electrostatic interactions\n', 'rcoulomb        = 1.0       ; Short-range electrostatic cut-off\n', 'rvdw            = 1.0       ; Short-range Van der Waals cut-off\n', 'pbc             = xyz       ; Periodic Boundary Conditions in all 3 dimensions\n']
mini = open("minim.mdp","w+")
mini.writelines(mini_list)
mini.close()
nvt_list = ['title                   =  NVT equilibration \n', 'define                  = -DPOSRES  ; position restrain the protein\n', '; Run parameters\n', 'integrator              = md        ; leap-frog integrator\n', 'nsteps                  = 50000     ; 2 * 50000 = 100 ps\n', 'dt                      = 0.002     ; 2 fs\n', '; Output control\n', 'nstxout                 = 500       ; save coordinates every 1.0 ps\n', 'nstvout                 = 500       ; save velocities every 1.0 ps\n', 'nstenergy               = 500       ; save energies every 1.0 ps\n', 'nstlog                  = 500       ; update log file every 1.0 ps\n', '; Bond parameters\n', 'continuation            = no        ; first dynamics run\n', 'constraint_algorithm    = lincs     ; holonomic constraints \n', 'constraints             = h-bonds   ; bonds involving H are constrained\n', 'lincs_iter              = 1         ; accuracy of LINCS\n', 'lincs_order             = 4         ; also related to accuracy\n', '; Nonbonded settings \n', 'cutoff-scheme           = Verlet    ; Buffered neighbor searching\n', 'ns_type                 = grid      ; search neighboring grid cells\n', 'nstlist                 = 10        ; 20 fs, largely irrelevant with Verlet\n', 'rcoulomb                = 1.0       ; short-range electrostatic cutoff (in nm)\n', 'rvdw                    = 1.0       ; short-range van der Waals cutoff (in nm)\n', 'DispCorr                = EnerPres  ; account for cut-off vdW scheme\n', '; Electrostatics\n', 'coulombtype             = PME       ; Particle Mesh Ewald for long-range electrostatics\n', 'pme_order               = 4         ; cubic interpolation\n', 'fourierspacing          = 0.16      ; grid spacing for FFT\n', '; Temperature coupling is on\n', 'tcoupl                  = V-rescale             ; modified Berendsen thermostat\n', 'tc-grps                 = Protein Non-Protein   ; two coupling groups - more accurate\n', 'tau_t                   = 0.1     0.1           ; time constant, in ps\n', 'ref_t                   = 300     300           ; reference temperature, one for each group, in K\n', '; Pressure coupling is off\n', 'pcoupl                  = no        ; no pressure coupling in NVT\n', '; Periodic boundary conditions\n', 'pbc                     = xyz       ; 3-D PBC\n', '; Velocity generation\n', 'gen_vel                 = yes       ; assign velocities from Maxwell distribution\n', 'gen_temp                = 300       ; temperature for Maxwell distribution\n', 'gen_seed                = -1        ; generate a random seed\n']
nvt = open("nvt.mdp","w+")
nvt.writelines(nvt_list)
nvt.close()
npt_list = ['title                   =  NPT equilibration \n', 'define                  = -DPOSRES  ; position restrain the protein\n', '; Run parameters\n', 'integrator              = md        ; leap-frog integrator\n', 'nsteps                  = 50000     ; 2 * 50000 = 100 ps\n', 'dt                      = 0.002     ; 2 fs\n', '; Output control\n', 'nstxout                 = 500       ; save coordinates every 1.0 ps\n', 'nstvout                 = 500       ; save velocities every 1.0 ps\n', 'nstenergy               = 500       ; save energies every 1.0 ps\n', 'nstlog                  = 500       ; update log file every 1.0 ps\n', '; Bond parameters\n', 'continuation            = yes       ; Restarting after NVT \n', 'constraint_algorithm    = lincs     ; holonomic constraints \n', 'constraints             = h-bonds   ; bonds involving H are constrained\n', 'lincs_iter              = 1         ; accuracy of LINCS\n', 'lincs_order             = 4         ; also related to accuracy\n', '; Nonbonded settings \n', 'cutoff-scheme           = Verlet    ; Buffered neighbor searching\n', 'ns_type                 = grid      ; search neighboring grid cells\n', 'nstlist                 = 10        ; 20 fs, largely irrelevant with Verlet scheme\n', 'rcoulomb                = 1.0       ; short-range electrostatic cutoff (in nm)\n', 'rvdw                    = 1.0       ; short-range van der Waals cutoff (in nm)\n', 'DispCorr                = EnerPres  ; account for cut-off vdW scheme\n', '; Electrostatics\n', 'coulombtype             = PME       ; Particle Mesh Ewald for long-range electrostatics\n', 'pme_order               = 4         ; cubic interpolation\n', 'fourierspacing          = 0.16      ; grid spacing for FFT\n', '; Temperature coupling is on\n', 'tcoupl                  = V-rescale             ; modified Berendsen thermostat\n', 'tc-grps                 = Protein Non-Protein   ; two coupling groups - more accurate\n', 'tau_t                   = 0.1     0.1           ; time constant, in ps\n', 'ref_t                   = 300     300           ; reference temperature, one for each group, in K\n', '; Pressure coupling is on\n', 'pcoupl                  = Berendsen             ; Pressure coupling on in NPT\n', 'pcoupltype              = isotropic             ; uniform scaling of box vectors\n', 'tau_p                   = 2.0                   ; time constant, in ps\n', 'ref_p                   = 1.0                   ; reference pressure, in bar\n', 'compressibility         = 4.5e-5                ; isothermal compressibility of water, bar^-1\n', 'refcoord_scaling        = com\n', '; Periodic boundary conditions\n', 'pbc                     = xyz       ; 3-D PBC\n', '; Velocity generation\n', 'gen_vel                 = no        ; Velocity generation is off \n']
npt = open("npt.mdp","w+")
npt.writelines(npt_list)
npt.close()
md_list = ['title                   =  MD \n', '; Run parameters\n', 'integrator              = md        ; leap-frog integrator\n', 'nsteps                  = {time}   ; 2 * 500000 = 1000 ps (1 ns)\n'.format(time = int(time*500000)), 'dt                      = 0.002     ; 2 fs\n', '; Output control\n', 'nstxout                 = 0         ; suppress bulky .trr file by specifying \n', 'nstvout                 = 0         ; 0 for output frequency of nstxout,\n', 'nstfout                 = 0         ; nstvout, and nstfout\n', 'nstenergy               = 5000      ; save energies every 10.0 ps\n', 'nstlog                  = 5000      ; update log file every 10.0 ps\n', 'nstxout-compressed      = 5000      ; save compressed coordinates every 10.0 ps\n', 'compressed-x-grps       = System    ; save the whole system\n', '; Bond parameters\n', 'continuation            = yes       ; Restarting after NPT \n', 'constraint_algorithm    = lincs     ; holonomic constraints \n', 'constraints             = h-bonds   ; bonds involving H are constrained\n', 'lincs_iter              = 1         ; accuracy of LINCS\n', 'lincs_order             = 4         ; also related to accuracy\n', '; Neighborsearching\n', 'cutoff-scheme           = Verlet    ; Buffered neighbor searching\n', 'ns_type                 = grid      ; search neighboring grid cells\n', 'nstlist                 = 10        ; 20 fs, largely irrelevant with Verlet scheme\n', 'rcoulomb                = 1.0       ; short-range electrostatic cutoff (in nm)\n', 'rvdw                    = 1.0       ; short-range van der Waals cutoff (in nm)\n', '; Electrostatics\n', 'coulombtype             = PME       ; Particle Mesh Ewald for long-range electrostatics\n', 'pme_order               = 4         ; cubic interpolation\n', 'fourierspacing          = 0.16      ; grid spacing for FFT\n', '; Temperature coupling is on\n', 'tcoupl                  = V-rescale             ; modified Berendsen thermostat\n', 'tc-grps                 = Protein Non-Protein   ; two coupling groups - more accurate\n', 'tau_t                   = 0.1     0.1           ; time constant, in ps\n', 'ref_t                   = 300     300           ; reference temperature, one for each group, in K\n', '; Pressure coupling is on\n', 'pcoupl                  = Parrinello-Rahman     ; Pressure coupling on in NPT\n', 'pcoupltype              = isotropic             ; uniform scaling of box vectors\n', 'tau_p                   = 2.0                   ; time constant, in ps\n', 'ref_p                   = 1.0                   ; reference pressure, in bar\n', 'compressibility         = 4.5e-5                ; isothermal compressibility of water, bar^-1\n', '; Periodic boundary conditions\n', 'pbc                     = xyz       ; 3-D PBC\n', '; Dispersion correction\n', 'DispCorr                = EnerPres  ; account for cut-off vdW scheme\n', '; Velocity generation\n', 'gen_vel                 = no        ; Velocity generation is off \n']
md = open("md.mdp","w+")
md.writelines(md_list)
md.close()

mmgbsa_DNA_list =['Sample input file for GB calculation\n', 'This input file is meant to show only that gmx_MMPBSA works. Althought,\n', 'we tried to used the input files as recommended in the Amber manual,\n', 'some parameters have been changed to perform more expensive calculations\n', 'in a reasonable amount of time. Feel free to change the parameters\n', 'according to what is better for your system.\n', '\n', '&general\n', 'sys_name="Prot-DNA"\n', 'startframe=5\n', 'endframe=14\n', 'PBRadii=4\n', 'forcefields="oldff/leaprc.ff99SB"\n', '/\n', '&gb\n', 'igb=8, saltcon=0.150, intdiel=10\n', '/\n', '&decomp\n', 'idecomp=2, dec_verbose=3,\n', 'print_res="within 4", csv_format=0,\n', '/']
mmgbsa_DNA = open("mmgbsa-DNA.in","w+")
mmgbsa_DNA.writelines(mmgbsa_DNA_list)
mmgbsa_DNA.close()

mmgbsa_PRO_list = ['Sample input file for GB calculation\n', 'This input file is meant to show only that gmx_MMPBSA works. Althought,\n', 'we tried to used the input files as recommended in the Amber manual,\n', 'some parameters have been changed to perform more expensive calculations\n', 'in a reasonable amount of time. Feel free to change the parameters\n', 'according to what is better for your system.\n', '\n', '&general\n', 'sys_name="Prot-Prot",\n', 'startframe=5,\n', 'endframe=14,\n', 'forcefields="oldff/leaprc.ff99SB",\n', '/\n', '&gb\n', 'igb=2, saltcon=0.150,\n', '/\n', '&decomp\n', 'idecomp=2, dec_verbose=3,\n', 'print_res="within 4", csv_format=0,\n', '/']
mmgbsa_PRO = open("mmgbsa-PRO.in","w+")
mmgbsa_PRO.writelines(mmgbsa_PRO_list)
mmgbsa_PRO.close()

mmgbsa_LIG_list = ['Sample input file for GB calculation\n', 'This input file is meant to show only that gmx_MMPBSA works. Althought,\n', 'we tried to used the input files as recommended in the Amber manual,\n', 'some parameters have been changed to perform more expensive calculations\n', 'in a reasonable amount of time. Feel free to change the parameters\n', 'according to what is better for your system.\n', '\n', '&general\n', 'sys_name="Prot-Lig-ST",\n', 'startframe=5,\n', 'endframe=14,\n', 'forcefields="oldff/leaprc.ff99SB,leaprc.gaff"\n', '/\n', '&gb\n', 'igb=5, saltcon=0.150,\n', '/\n', '&decomp\n', 'idecomp=2, dec_verbose=3,\n', 'print_res="within 4", csv_format=0,\n', '/']
mmgbsa_LIG = open("mmgbsa-LIG.in","w+")
mmgbsa_LIG.writelines(mmgbsa_LIG_list)
mmgbsa_LIG.close()

mmgbsa_METAL_list =['Sample input file for GB calculation\n', 'This input file is meant to show only that gmx_MMPBSA works. Althought,\n', 'we tried to used the input files as recommended in the Amber manual,\n', 'some parameters have been changed to perform more expensive calculations\n', 'in a reasonable amount of time. Feel free to change the parameters\n', 'according to what is better for your system.\n', '\n', '&general\n', 'sys_name="Metalloprotein-peptide",\n', 'startframe=5\n', 'endframe=14,\n', 'ions_parameters=1,\n', '/\n', '&gb\n', 'igb=2, saltcon=0.150,\n', '/\n', '\n', '&decomp\n', 'idecomp=2, dec_verbose=3,\n', 'print_res="within 4", csv_format=0,\n', '/']
mmgbsa_METAL = open("mmgbsa-METAL.in","w+")
mmgbsa_METAL.writelines(mmgbsa_METAL_list)
mmgbsa_METAL.close()

LandScape_list = ["# -*- coding: UTF-8 -*-\n",'import os\n', 'print("protein-protein")\n', 'os.system("gmx trjconv -s md_0_1.tpr -f md_0_1_noPBC_MOL.xtc -o md_0_1_noPBC_fit.xtc -fit rot+trans")\n', 'print("c-alpha-c-alpha")\n', 'os.system("gmx covar -s md_0_1.tpr -f md_0_1_noPBC_fit.xtc -o eigenvalues.xvg -v eigenvectors.trr -xpma covapic.xpm ")\n', 'print("c-alpha-c-alpha")\n', 'os.system("gmx anaeig -s md_0_1.tpr -f md_0_1_noPBC_fit.xtc -v eigenvectors.trr -first 1 -last 1 -proj pc1.xvg ")\n', 'print("c-alpha-c-alpha")\n', 'os.system("gmx anaeig -s md_0_1.tpr -f md_0_1_noPBC_fit.xtc -v eigenvectors.trr -first 2 -last 2 -proj pc2.xvg  ")\n', '\n', 'pc1 = open("pc1.xvg","r")\n', 'pc1_lines = [i for i in pc1]\n', 'pc1.close()\n', '\n', 'pc2 = open("pc2.xvg","r")\n', 'pc2_lines = [i for i in pc2]\n', 'pc2.close()\n', '\n', 'head = []\n', 'pc1_list = []\n', 'pc2_list = []\n', 'for i in range(len(pc1_lines)):\n', '    if pc1_lines[i][0] == "@":\n', '        head.append(pc1_lines[i])\n', '    else:\n', '        pc1_list.append(pc1_lines[i][0:-2])\n', '# print(pc1_list)\n', 'for i in range(len(pc2_lines)):\n', '    if pc2_lines[i][0] != "@":\n', '        pc2_list.append(pc2_lines[i][-10:-1])\n', '# print(pc2_list)\n', 'for i in range(len(pc1_list)):\n', '    pc1_list[i] = pc1_list[i] + "  " + pc2_list[i] + "\\n"\n', '\n', 'print(pc1_list)\n', '\n', 'pc12 = open("pc12_sham.xvg","w+")\n', 'pc12.writelines(head+pc1_list)\n', 'pc12.close()\n', '\n', 'os.system("gmx sham -tsham 300 -nlevels 100 -f pc12_sham.xvg -ls pc12_gibbs.xpm -g pc_12.log -lsh pc12_enthalpy.xpm -lss pc12_entropy.xpm")\n', '\n', '# os.system("pip install DuIvyTools  -i https://pypi.tuna.tsinghua.edu.cn/simple")\n', '# os.system("dit xpm_show -f pc12_gibbs.xpm -ip")\n', '# os.system("dit xpm_show -f pc12_gibbs.xpm -3d")']
LandScape = open("LandScape-PCA.py","w+")
LandScape.writelines(LandScape_list)
LandScape.close()

LandScape_list2 = ["# -*- coding: UTF-8 -*-\n",'import os\n', '# print("protein-protein")\n', 'os.system("gmx trjconv -s md_0_1.tpr -f md_0_1_noPBC_MOL.xtc -o md_0_1_noPBC_fit.xtc -fit rot+trans")\n', '# 选择backbone进行计算和输出\n', 'os.system("gmx rms -s md_0_1.tpr -f md_0_1_noPBC_fit.xtc -o rmsd-landscape.xvg")\n', '# 同样选择backbone\n', 'os.system("gmx gyrate -s md_0_1.tpr -f md_0_1_noPBC_fit.xtc -o gyrate-landscape.xvg")\n', '\n', '\n', 'rmsd = open("rmsd-landscape.xvg","r")\n', 'rmsd_lines = [i for i in rmsd]\n', 'rmsd.close()\n', '\n', 'gyrate = open("gyrate-landscape.xvg","r")\n', 'gyrate_lines = [i for i in gyrate]\n', 'gyrate.close()\n', '\n', '\n', 'head = []\n', 'pc1_list = []\n', 'pc2_list = []\n', 'for i in range(len(rmsd_lines)):\n', '    if rmsd_lines[i][0] == "@":\n', '        head.append(rmsd_lines[i])\n', '    else:\n', '        pc1_list.append(rmsd_lines[i][0:-2])\n', '# print(pc1_list)\n', 'for i in range(len(gyrate_lines)):\n', '    if gyrate_lines[i][0] != "@":\n', '        pc2_list.append(gyrate_lines[i][15:22])\n', '# print(pc2_list)\n', 'for i in range(len(pc1_list)):\n', '    pc1_list[i] = pc1_list[i] + "    " + pc2_list[i] + "\\n"\n', '\n', '# print(pc1_list)\n', '\n', 'pc12 = open("rmsd_gyrate_sham.xvg","w+")\n', 'pc12.writelines(head+pc1_list)\n', 'pc12.close()\n', '\n', '\n', 'os.system("gmx sham -tsham 310 -nlevels 100 -f rmsd_gyrate_sham.xvg -ls rmsd_gyrate_gibbs.xpm -g rmsd_gyrate.log -lsh enthalpy.xpm -lss entropy.xpm")\n', '#os.system("gmx sham -tsham 300 -nlevels 100 -f pc12_sham.xvg -ls pc12_gibbs.xpm -g pc_12.log -lsh pc12_enthalpy.xpm -lss pc12_entropy.xpm")\n', '\n', '# os.system("pip install DuIvyTools  -i https://pypi.tuna.tsinghua.edu.cn/simple")\n', '# os.system("dit xpm_show -f pc12_gibbs.xpm -ip -x RMSD -y Gyrate")\n', '# os.system("dit xpm_show -f pc12_gibbs.xpm -3d -x RMSD -y Gyrate")\n','os.system("dit xpm_show -f rmsd_gyrate_gibbs.xpm -ip -x RMSD -y Gyrate -o LandScape-2d.tif -ns")\n','os.system("dit xpm_show -f rmsd_gyrate_gibbs.xpm -3d -x RMSD -y Gyrate -o LandScape-3d.tif -ns")']
LandScape2 = open("LandScape-RMSD-Gyrate.py","w+")
LandScape2.writelines(LandScape_list2)
LandScape2.close()

Continue_list = ["# -*- coding: UTF-8 -*-\n",'import os\n', 'import sys\n', '\n',
               'os.system("gmx mdrun -s md_0_1.tpr -cpi md_0_1.cpt -deffnm md_0_1 -update gpu -v")\n']
Continue = open("Continue_RunMD.py","w+")
Continue.writelines(Continue_list)
Continue.close()

#运行md
# os.system("gmx pdb2gmx -f {} -o Pro_temp_H.gro -ff amber99sb-ildn -water tip3p -ignh".format(name))
os.system("gmx editconf -f Pro_temp_H.gro -o Pro_temp_H_box.gro -c -d 1")
os.system("gmx solvate -cp Pro_temp_H_box.gro -o Pro_temp_H_box_water.gro -p topol.top")
os.system("gmx grompp -f ions.mdp -c Pro_temp_H_box_water.gro -p topol.top -o ions.tpr -maxwarn 99")
os.system("gmx genion -s ions.tpr -o Pro_temp_H_box_water_ion.gro -p topol.top -neutral")
#os.system("screen -L -dmS RunMD python RunMD.py")


os.system("gmx grompp -f minim.mdp -c Pro_temp_H_box_water_ion.gro -p topol.top -o em.tpr -maxwarn 99")
os.system("gmx mdrun -v -deffnm em")
os.system("gmx grompp -f nvt.mdp -c em.gro -r em.gro -p topol.top -o nvt.tpr -maxwarn 99")
os.system("gmx mdrun -deffnm nvt -update gpu -v")
os.system("gmx grompp -f npt.mdp -c nvt.gro -r nvt.gro -t nvt.cpt -p topol.top -o npt.tpr -maxwarn 99")
os.system("gmx mdrun -deffnm npt -update gpu -v")
os.system("gmx grompp -f md.mdp -c npt.gro -t npt.cpt -p topol.top -o md_0_1.tpr -maxwarn 99")
os.system("gmx mdrun -deffnm md_0_1 -update gpu -v")




